from dotenv import load_dotenv
import requests
import os

from langchain_core.output_parsers import JsonOutputKeyToolsParser, StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnableLambda
from langchain_core.tools import tool
from langchain_ollama import ChatOllama
from pydantic import BaseModel, Field

load_dotenv(override=True)


class WeatherQuery(BaseModel):
    loc: str = Field(description="城市名称")


@tool(args_schema=WeatherQuery)
def get_weather(loc: str) -> str:
    """
    获取指定城市的天气信息
    :param loc:必要参数，字符串类型，用于指定城市的名称，注意中国的城市需要用对应城市的英文名称代替换。
    比如北京需要用"Beijing"代替。
    """
    API_KEY = os.getenv("OPENWEATHERMAP_API_KEY")
    if not API_KEY:
        return "未设置 API 密钥"

    url = f"http://api.openweathermap.org/data/2.5/weather?q={loc}&appid={API_KEY}&units=metric&lang=zh_cn"
    try:
        print(url)
        response = requests.get(url)
        data = response.json()

        if data.get("cod") != 200:
            return f"查询失败: {data.get('message')}"

        # weather = data["weather"][0]["description"]
        # temp = data["main"]["temp"]
        # humidity = data["main"]["humidity"]
        return data
    except Exception as e:
        return f"请求失败: {str(e)}"


if __name__ == "__main__":
    DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
    # model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    model = ChatOllama(model="qwen3:latest", base_url="http://192.168.97.217:11434")
    parser = JsonOutputKeyToolsParser(key_name=get_weather.name, firt_tool_only=True)
    tools = [get_weather]
    llm_with_tools = model.bind_tools(tools)
    extract_dict = RunnableLambda(lambda x: x[0])
    weather_chain_1 = llm_with_tools | parser | extract_dict | get_weather
    response1 = weather_chain_1.invoke("需要将中文的城市为国际标识，请告诉我北京的天气情况")
    print(response1)
    output_prompt = PromptTemplate.from_template("""
        你将收到一段 JSON格式的天气数据,请用简洁自然的方式将其转述给用户。
        一下是天气JSON数据:
        ``` json
        {weather_json}
        ```
        请将其转为中文描述，例如：
        “北京天气晴，气温为23℃,湿度58%，风速5.0米/秒”
        “只返回一句话描述，不要其他说明或解释”
    """)
    out_chain = output_prompt | model | StrOutputParser()
    response2 = out_chain.invoke({"weather_json": response1})
    print(response2)
